/*
 * Wazuh Vulnerability scanner
 * Copyright (C) 2015, Wazuh Inc.
 * Nov 13, 2023.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#include "argsParser.hpp"
#include "flatbuffers/flatbuffers.h"
#include "flatbuffers/idl.h"
#include "flatbuffers/include/messageBuffer_generated.h"
#include "flatbuffers/include/rsync_generated.h"
#include "flatbuffers/include/syscollector_deltas_generated.h"
#include "rocksDBWrapper.hpp"
#include <iostream>

namespace Log
{
    std::function<void(
        const int, const std::string&, const std::string&, const int, const std::string&, const std::string&, va_list)>
        GLOBAL_LOG_FUNCTION;
}; // namespace Log

int main(const int argc, const char** argv)
{
    try
    {
        CmdLineArgs cmdLineArgs(argc, argv);

        auto rocksDB = Utils::RocksDBWrapper(cmdLineArgs.getDBPath());
        auto fbs = cmdLineArgs.getFbsPath();
        auto requestedKey = cmdLineArgs.getKey();
        auto seekKey = cmdLineArgs.getSeekKey();
        auto columnFamily = cmdLineArgs.getColumnFamily();
        auto requestedValue = cmdLineArgs.getValue();
        auto isMessageBuffer = cmdLineArgs.getValueIsMessageBuffer();
        auto fbsDelta = cmdLineArgs.getFbsDeltaPath();
        auto fbsSync = cmdLineArgs.getFbsSyncPath();

        flatbuffers::IDLOptions options;
        options.strict_json = true;
        flatbuffers::Parser fbParser(options);
        flatbuffers::Parser deltaParser(options);
        flatbuffers::Parser syncParser(options);

        auto parseSchema = [&](const std::string& fbsPath, flatbuffers::Parser& parser)
        {
            std::string schemaStr;
            if (!fbsPath.empty())
            {
                if (!flatbuffers::LoadFile(fbsPath.c_str(), false, &schemaStr))
                {
                    throw std::runtime_error("Unable to load schema file: " + fbsPath);
                }
                if (!parser.Parse(schemaStr.c_str()))
                {
                    throw std::runtime_error("Unable to parse schema file: " + fbsPath);
                }
            }
        };

        parseSchema(fbs, fbParser);
        parseSchema(fbsDelta, deltaParser);
        parseSchema(fbsSync, syncParser);

        auto printValue = [&](const std::string& key, const auto& slice)
        {
            if (isMessageBuffer.empty())
            {
                if (!fbs.empty())
                {
                    std::string strData;
                    flatbuffers::GenText(fbParser, reinterpret_cast<const uint8_t*>(slice.data()), &strData);
                    std::cout << key << " ==> " << strData << std::endl;
                }
                else
                {
                    std::cout << key << " ==> " << slice.ToString() << std::endl;
                }
            }
            else
            {
                auto message = GetMessageBuffer(slice.data());

                if (message->type() == BufferType::BufferType_RSync)
                {
                    std::string strData;
                    flatbuffers::GenText(syncParser, reinterpret_cast<const uint8_t*>(message->data()), &strData);
                    std::cout << key << " ==> " << strData << std::endl;
                }
                else if (message->type() == BufferType::BufferType_DBSync)
                {
                    std::string strData;
                    flatbuffers::GenText(deltaParser, reinterpret_cast<const uint8_t*>(message->data()), &strData);
                    std::cout << key << " ==> " << strData << std::endl;
                }
                else if (message->type() == BufferType::BufferType_JSON)
                {
                    try
                    {
                        auto jsonData = nlohmann::json::parse(message->data()->data(),
                                                              message->data()->data() + message->data()->size());

                        std::cout << key << " ==> " << jsonData.dump() << std::endl;
                    }
                    catch (const std::exception& e)
                    {
                        std::cout << "ERROR. Unable to parse key: '" << key << "'. JSON: '" << message->data()->data()
                                  << "'. Reason: " << e.what() << std::endl;
                    }
                }
                else
                {
                    throw std::runtime_error("Unknown event type");
                }
            }
        };

        if (!seekKey.empty())
        {
            for (const auto& [key, value] : rocksDB.seek(seekKey, columnFamily))
            {
                printValue(key, value);
            }
        }
        else if (!requestedValue.empty() && !requestedKey.empty())
        {
            if (!rocksDB.columnExists(columnFamily))
            {
                rocksDB.createColumn(columnFamily);
            }

            if (!fbs.empty())
            {
                if (!fbParser.Parse(requestedValue.c_str()))
                {
                    throw std::runtime_error("Unable to parse value.");
                }
                rocksdb::Slice flatbufferResource(reinterpret_cast<const char*>(fbParser.builder_.GetBufferPointer()),
                                                  fbParser.builder_.GetSize());
                rocksDB.put(requestedKey, flatbufferResource, columnFamily);
            }
            else
            {
                rocksDB.put(requestedKey, requestedValue, columnFamily);
            }

            std::cout << "Value inserted." << std::endl;
        }
        else if (!requestedKey.empty())
        {
            rocksdb::PinnableSlice slice;
            if (!rocksDB.get(requestedKey, slice, columnFamily))
            {
                throw std::runtime_error("Unable to find resource.");
            }
            printValue(requestedKey, slice);
        }
        else
        {
            for (const auto& [key, value] : rocksDB.begin(columnFamily))
            {
                printValue(key, value);
            }
        }
    }
    catch (const std::exception& e)
    {
        std::cerr << e.what() << std::endl;
        CmdLineArgs::showHelp();
        return 1;
    }

    return 0;
}
